Codecov Report

Merging #7032 (6d8fff3) into master (7ca4173) will decrease coverage by 4%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #7032    +/-   ##
=======================================
- Coverage      92%     88%    -4%     
=======================================
  Files         196     196            
  Lines       12835   12833     -2     
=======================================
- Hits        11832   11268   -564     
- Misses       1003    1565   +562     

Can you write a test that fails on current master without the fix?

I'm sorry that I have no idea how to write a test for this issue...

Even if update_learning_rates is not called, training itself will proceed. Also, I can't think of a good way to check that update_learning_rates is not being called, from outside the Trainer.

@tkng You must have observed this bug somehow in your code right? If you provide a minimal example (see our bug_report_model), we can easily help you convert it to a good regression test :)

@awaelchli thanks, I wrote a minimal example and uploaded to https://gist.github.com/tkng/a89ffdab2558aedc04a67f7142e6aeb0. (Sorry, I noticed to bug_report_model.py after I wrote this example, so the code is based on lightning-flash.)

Once you run this example and opened the log directory by tensorboard, you'll see this learning rate decay. Pink one is with current pytorch-lightning, green one is with the patch attached on this PR.

Thanks. Based on your description and code, I pushed a minimal test that fails on master.

Why is the learning rate update dependent on the validation loop at all?

I suspect because of this scheduler
https://pytorch.org/docs/stable/optim.html#torch.optim.lr_scheduler.ReduceLROnPlateau
But I'm not an expert on this section of the code. I believe @rohitgr7 has worked a bit here if I'm not mistaken.

well, the solution here is a bit incorrect. val_loop_called actually tells whether the val_loop is called within the training loop or not for the cases where val_check_interval(float) < 1. or val_check_interval < num_train_batches. In such a case val_loop_called will be True and thus we need to ensure to explicitly call the epoch level schedulers explicitly outside the train loop if there is no epoch end validation going to happen. I have updated a test #7084 which will fail with your changes.

Epoch level schedulers also update within the evaluation loop because we need to ensure they get updated after the validation loop and right before checkpointing to ensure the correct scheduler state dict being saved in the checkpoint file. I know this is a bit complicated but most of the complexity is due to this ReduceLROnPlateau scheduler which needs a monitor to update itself. This monitor can belong to anything logged in either train or val.

To handle the cases where check_val_every_n_epoch > 1, I think we need to somehow pull out the scheduler update from val_loop and still need to ensure that correct scheduler state_dict gets updated within the checkpoint file (we have tests for that) and ofcourse should work with ReduceLROnPlateau (I think we have tests for that too).

@rohitgr7 Hmm, the problem seems more complicated than I had thought...

May I assume that update_learning_rates should always be called at the end of one training epoch, even for ReduceLROnPlateau?

If the assumption is correct, then I think the correct code would be to call update_learning_rate immediately after the end of the training process for one epoch, in the run_training_epoch method of TrainLoop class.

I'm thinking this understanding is consistent with your statement, I think we need to somehow pull out the scheduler update from val_loop and still need to ensure that correct scheduler state_dict gets updated within the checkpoint file.

I tried to write such a code once, but I could not make the unit test pass. On the other hand, if I discarded the above understanding and just considered passing the unit test, I only needed to change one line, as I just pushed. (Just for sure, this also passes the test for #7084.) However, I am not sure whether my current code is correct. 🤔

yeap val_loop_called is redundant since should_check_val can cover up that case. Calling lr_update right after the train loop will create problem with ReduceLROnPlateau since the monitor might not be logged by then. Also even if we do it after the validation loop then checkpoint won't be getting the correct lr_state_dict. I think I have a solution for this. Will try it on weekend.

Converting to draft so merging is blocked until these comments are addressed 👍

@tkng how is it going here, still WIP or ready to go? 🐰

@Borda This pull request passes the current all unit tests, but according to @rohitgr7, it is incompatible with ReduceLROnPlateau.

I have tried to rewrite this PR to compatible with ReduceLROnPlateau, but I have failed to write such code that passes all tests and still conforms to ReduceLROnPlateau. I think we have two options.

1. wait for new code from @rohitgr7
2. merge this pull request, since it passes all the current unit tests

I am OK with either, but we'll need a opinion from @rohitgr7.

I tried something, but couldn't find an optimal solution that can cover all the cases because of this ReduceLROnPlateau. The problem is if a monitor is used with this scheduler with check_val_every_n_epoch > 1 then there is no way to decide whether to do a scheduler.step after validation or after every train loop because this monitor can be anything and can be logged anywhere. For eg check_val_every_n_epoch>1, if it's logged in val loop then we ideally shouldn't do scheduler.step after every epoch but if it's logged in train loop then we should. Although one can easily configure this correctly by setting scheduler['frequency'] accordingly if they understand how they have configured their train/val loop using trainer parameters.

So I'd say merge this one.